package com.pudding.ragdemo;

import com.pudding.tools.ModelUtil;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.injector.ContentInjector;
import dev.langchain4j.rag.content.injector.DefaultContentInjector;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.redis.RedisEmbeddingStore;

import java.util.List;

public class MeituanRagDemo {
    public static void main(String[] args) {
        // 1. 将问题进行向量化，并且去redis中查询相似度比较高的知识条目
        // 创建向量模型
        EmbeddingModel embeddingModel = OpenAiEmbeddingModel.builder()
                .baseUrl(ModelUtil.BASE_URL_OPENAI)
                .apiKey(ModelUtil.API_KEY_OPENAI)
                .build();
        // 创建向量存储模型
        EmbeddingStore<TextSegment> embeddingStore = RedisEmbeddingStore.builder()
                .host("127.0.0.1")
                .port(6379)
                .dimension(1536)
                .indexName("meituan-rag")
                .build();
        // 创建内容检索模型，无需关注具体的存储，这个是上层的抽象
        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore) //向量存储模型
                .embeddingModel(embeddingModel) //向量模型
                .maxResults(5) // 最相似的5个结果
                .minScore(0.8) // 只找相似度在0.8以上的内容
                .build();

        String question = "在线⽀付取消订单后钱怎么返还？"; //⽤户的问题

        Query query = new Query(question);
        List<Content> contentList = contentRetriever.retrieve(query); // 检索内容

        // 打印相关信息
        for (Content content : contentList) {
            System.out.println(content);
        }

        // 2. 将问题和相关知识条目整合成一个完整的信息，发送给大模型
        ContentInjector contentInjector = new DefaultContentInjector(); // 内容注入器, 将内容和聊天信息整合到一起，顶层接口，我们可以自己实现
        UserMessage promptMessage = contentInjector.inject(contentList, UserMessage.from(question));
        for (dev.langchain4j.data.message.Content content : promptMessage.contents()) {
            System.out.println(content);
        }
        System.out.println("-------------------------------");
        // 创建模型,将问题和内容整合到一起，发送给模型
        ChatLanguageModel model = OpenAiChatModel.builder()
                .apiKey(ModelUtil.API_KEY_OPENAI)
                .build();
        Response<AiMessage> generate = model.generate(promptMessage);
        System.out.println(generate.content().text());
    }
}
